#!/usr/bin/env python3
"""
ArUco标记跟踪器类

该模块实现了 ArucoTracker 类，提供基于海康威视相机的ArUco标记实时跟踪功能。
"""
import cv2
import numpy as np
import json
import math
import time
import threading
from typing import Dict, List, Optional, Callable, Tuple, Any

try:
    from hik_camera import HikCamera
except ImportError as e:
    print(f"警告: 导入 hik_camera 失败: {e}")
    print("请确保 hik_camera 包在正确的路径下")
    HikCamera = None

from .marker import ArucoMarker


class ArucoTracker:
    """
    ArUco标记跟踪器
    
    基于海康威视相机的ArUco标记实时跟踪系统，支持多标记同时跟踪、
    状态管理、相对位置计算等功能。
    """
    
    def __init__(self, camera_index: int = 0, calibration_file: Optional[str] = None,
                 max_lost_frames: int = 30, aruco_dict=cv2.aruco.DICT_6X6_50):
        """
        初始化ArUco跟踪器
        
        Args:
            camera_index (int): 相机索引
            calibration_file (str): 相机标定参数文件路径
            max_lost_frames (int): 最大允许丢失的帧数
            aruco_dict: ArUco字典类型
        """
        self.camera_index = camera_index
        self.calibration_file = calibration_file
        self.max_lost_frames = max_lost_frames
        
        # 相机相关
        self.camera = None
        self.camera_matrix = None
        self.dist_coeffs = None
        self.calibrated = False
        
        # ArUco检测器
        self.aruco_dict = cv2.aruco.getPredefinedDictionary(aruco_dict)
        self.aruco_params = cv2.aruco.DetectorParameters()
        self.cv2_detector = cv2.aruco.ArucoDetector(self.aruco_dict, self.aruco_params)
        
        # 标记管理
        self.registered_markers: Dict[str, ArucoMarker] = {}  # name -> ArucoMarker
        self.marker_id_map: Dict[int, str] = {}  # marker_id -> name
        
        # 跟踪状态
        self.is_tracking = False
        self.tracking_thread = None
        self.stop_event = threading.Event()
        self._lock = threading.Lock()  # 线程锁，保护共享数据
        
        # 回调函数
        self.user_callback: Optional[Callable] = None
        
        # 统计信息
        self.frame_count = 0
        self.detection_count = 0
        self.fps = 0.0
        self.last_fps_time = time.time()
        
        # 加载标定参数
        if calibration_file:
            self.load_calibration(calibration_file)
        
        print(f"ArucoTracker 初始化完成 (相机索引: {camera_index})")
    
    def load_calibration(self, filename: str) -> bool:
        """
        加载相机标定参数
        
        Args:
            filename (str): 标定参数文件路径
            
        Returns:
            bool: 是否成功加载
        """
        try:
            if filename.endswith('.npz'):
                data = np.load(filename)
                self.camera_matrix = data['camera_matrix']
                self.dist_coeffs = data['dist_coeffs']
            elif filename.endswith('.json'):
                with open(filename, 'r') as f:
                    data = json.load(f)
                self.camera_matrix = np.array(data['camera_matrix'])
                self.dist_coeffs = np.array(data['distortion_coefficients'])
            else:
                raise ValueError("不支持的文件格式，请使用 .npz 或 .json 文件")
            
            self.calibrated = True
            print(f"标定参数加载成功: {filename}")
            return True
            
        except Exception as e:
            print(f"加载标定参数失败: {e}")
            return False
    
    def register_marker(self, marker: ArucoMarker) -> bool:
        """
        注册ArUco标记（线程安全）
        
        Args:
            marker (ArucoMarker): 要注册的标记对象
            
        Returns:
            bool: 是否成功注册
        """
        with self._lock:
            if marker.name in self.registered_markers:
                print(f"警告: 标记名称 '{marker.name}' 已存在")
                return False
            
            if marker.marker_id in self.marker_id_map:
                existing_name = self.marker_id_map[marker.marker_id]
                print(f"警告: 标记ID {marker.marker_id} 已被 '{existing_name}' 使用")
                return False
            
            self.registered_markers[marker.name] = marker
            self.marker_id_map[marker.marker_id] = marker.name
            
            print(f"成功注册标记: {marker.name} (ID: {marker.marker_id})")
            return True
    
    def unregister_marker(self, name: str) -> bool:
        """
        注销ArUco标记（线程安全）
        
        Args:
            name (str): 标记名称
            
        Returns:
            bool: 是否成功注销
        """
        with self._lock:
            if name not in self.registered_markers:
                print(f"警告: 标记 '{name}' 不存在")
                return False
            
            marker = self.registered_markers[name]
            del self.registered_markers[name]
            del self.marker_id_map[marker.marker_id]
            
            print(f"成功注销标记: {name}")
            return True
    
    def set_callback(self, callback: Callable) -> None:
        """
        设置用户回调函数
        
        Args:
            callback (Callable): 回调函数，接收 tracker 实例作为参数
        """
        self.user_callback = callback
        print("用户回调函数已设置")
    
    def start_tracking(self) -> bool:
        """
        开始跟踪
        
        Returns:
            bool: 是否成功开始跟踪
        """
        if self.is_tracking:
            print("跟踪已在运行中")
            return True
        
        if HikCamera is None:
            print("错误: HikCamera 类不可用")
            return False
        
        # 初始化相机
        try:
            self.camera = HikCamera(camera_index=self.camera_index)
            print("相机初始化成功")
        except Exception as e:
            print(f"相机初始化失败: {e}")
            return False
        
        # 启动跟踪线程
        self.is_tracking = True
        self.stop_event.clear()
        self.tracking_thread = threading.Thread(target=self._tracking_loop, daemon=True)
        self.tracking_thread.start()
        
        print("开始ArUco标记跟踪...")
        return True
    
    def stop_tracking(self) -> None:
        """停止跟踪"""
        if not self.is_tracking:
            print("跟踪未在运行")
            return
        
        print("停止ArUco标记跟踪...")
        
        self.is_tracking = False
        self.stop_event.set()
        
        # 等待跟踪线程结束
        if self.tracking_thread is not None and self.tracking_thread.is_alive():
            self.tracking_thread.join(timeout=2.0)
            if self.tracking_thread.is_alive():
                print("警告: 跟踪线程未能正常结束")
        
        # 关闭相机
        if self.camera:
            self.camera.close()
            self.camera = None
        
        print("跟踪已停止")
    
    def _tracking_loop(self) -> None:
        """
        跟踪循环，在子线程中运行
        """
        print("跟踪线程已启动")
        
        try:
            while self.is_tracking and not self.stop_event.is_set():
                # 从相机读取一帧图像
                self.camera.read_frame()
                frame = self.camera.frame
                
                if frame is not None:
                    # 处理图像帧
                    self._process_frame(frame)
                else:
                    # 如果没有获取到图像，短暂休眠
                    time.sleep(0.001)
                    
        except Exception as e:
            print(f"跟踪循环错误: {e}")
        finally:
            print("跟踪线程已结束")
    
    def _process_frame(self, frame: np.ndarray) -> None:
        """
        处理图像帧
        
        Args:
            frame (np.ndarray): 图像帧
        """
        try:
            # 更新帧统计
            self.frame_count += 1
            current_time = time.time()
            
            # 计算FPS
            if current_time - self.last_fps_time >= 1.0:
                self.fps = self.frame_count / (current_time - self.last_fps_time)
                self.frame_count = 0
                self.last_fps_time = current_time
            
            # 图像去畸变
            if self.calibrated:
                frame = self._undistort_image(frame)
            
            # 检测ArUco标记
            detected_markers = self._detect_aruco_markers(frame)
            
            # 更新已注册标记的状态
            self._update_marker_states(detected_markers)
            
            # 调用用户回调函数
            if self.user_callback:
                try:
                    self.user_callback(self)
                except Exception as e:
                    print(f"用户回调函数执行错误: {e}")
                    
        except Exception as e:
            print(f"处理图像帧错误: {e}")
        
    def _camera_callback(self, frame: np.ndarray, frame_info, camera_instance) -> None:
        """
        相机回调函数（已弃用，保留以兼容旧代码）
        
        Args:
            frame (np.ndarray): 图像帧
            frame_info: 帧信息
            camera_instance: 相机实例
        """
        print("警告: 正在使用已弃用的回调函数模式，请使用新的线程模式")
        if not self.is_tracking or frame is None:
            return
        
        self._process_frame(frame)
    
    def _undistort_image(self, image: np.ndarray) -> np.ndarray:
        """
        图像去畸变
        
        Args:
            image (np.ndarray): 输入图像
            
        Returns:
            np.ndarray: 去畸变后的图像
        """
        if not self.calibrated:
            return image
        
        h, w = image.shape[:2]
        newcameramtx, roi = cv2.getOptimalNewCameraMatrix(
            self.camera_matrix, self.dist_coeffs, (w, h), 1, (w, h)
        )
        
        undistorted = cv2.undistort(image, self.camera_matrix, self.dist_coeffs, None, newcameramtx)
        return undistorted
    
    def _detect_aruco_markers(self, image: np.ndarray) -> List[Dict]:
        """
        检测ArUco标记
        
        Args:
            image (np.ndarray): 输入图像
            
        Returns:
            List[Dict]: 检测到的标记信息列表
        """
        # 检测标记

        corners, ids, rejected = self.cv2_detector.detectMarkers(image)
        
        detected_markers = []
        
        if ids is not None and self.calibrated:
            for i, marker_id in enumerate(ids.flatten()):
                # 检查是否为已注册的标记
                if marker_id not in self.marker_id_map:
                    continue
                
                marker_name = self.marker_id_map[marker_id]
                registered_marker = self.registered_markers[marker_name]
                
                # 估计姿态
                # 使用 solvePnP 估计姿态
                obj_points = np.array([
                    [0, 0, 0],
                    [registered_marker.size_mm, 0, 0],
                    [registered_marker.size_mm, registered_marker.size_mm, 0],
                    [0, registered_marker.size_mm, 0]
                ], dtype=np.float32)
                img_points = corners[i].reshape(-1, 2).astype(np.float32)
                success, rvec, tvec = cv2.solvePnP(
                    obj_points, img_points,
                    self.camera_matrix, self.dist_coeffs
                )
                if not success:
                    continue
                rvecs = [rvec]
                tvecs = [tvec]
                
                # 转换旋转向量为欧拉角
                rotation_matrix, _ = cv2.Rodrigues(rvecs[0])
                roll, pitch, yaw = self._rotation_matrix_to_euler(rotation_matrix)
                
                # 计算距离
                distance = np.linalg.norm(tvecs[0])
                
                detected_markers.append({
                    'name': marker_name,
                    'marker_id': marker_id,
                    'position': tvecs[0].flatten(),
                    'rotation': np.array([roll, pitch, yaw]),
                    'corners': corners[i],
                    'distance': distance
                })
        
        return detected_markers
    
    def _rotation_matrix_to_euler(self, R: np.ndarray) -> Tuple[float, float, float]:
        """
        将旋转矩阵转换为欧拉角
        
        Args:
            R (np.ndarray): 3x3旋转矩阵
            
        Returns:
            Tuple[float, float, float]: (roll, pitch, yaw) 欧拉角(度)
        """
        sy = math.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
        singular = sy < 1e-6
        
        if not singular:
            x = math.atan2(R[2, 1], R[2, 2])
            y = math.atan2(-R[2, 0], sy)
            z = math.atan2(R[1, 0], R[0, 0])
        else:
            x = math.atan2(-R[1, 2], R[1, 1])
            y = math.atan2(-R[2, 0], sy)
            z = 0
        
        return math.degrees(x), math.degrees(y), math.degrees(z)
    
    def _update_marker_states(self, detected_markers: List[Dict]) -> None:
        """
        更新标记状态（线程安全）
        
        Args:
            detected_markers (List[Dict]): 检测到的标记列表
        """
        with self._lock:
            # 创建检测到的标记名称集合
            detected_names = {marker['name'] for marker in detected_markers}
            
            # 更新检测到的标记
            for marker_info in detected_markers:
                marker = self.registered_markers[marker_info['name']]
                marker.update_detection(
                    marker_info['position'],
                    marker_info['rotation'],
                    marker_info['corners'],
                    marker_info['distance']
                )
                self.detection_count += 1
            
            # 更新未检测到的标记
            for name, marker in self.registered_markers.items():
                if name not in detected_names:
                    marker.update_lost()
                    
                    # 如果丢失帧数超过阈值，标记为丢失
                    if marker.lost_frames >= self.max_lost_frames and marker.is_tracked:
                        marker.mark_lost()
    
    def get_marker_status(self, name: str) -> Optional[Dict]:
        """
        获取标记状态（线程安全）
        
        Args:
            name (str): 标记名称
            
        Returns:
            Dict or None: 标记状态信息，如果标记不存在返回None
        """
        with self._lock:
            if name not in self.registered_markers:
                print(f"标记 '{name}' 不存在")
                return None
            
            marker = self.registered_markers[name]
            
            if not marker.is_tracked:
                print(f"标记 '{name}' 当前未被跟踪")
            
            return marker.get_status()
    
    def get_relative_position(self, name1: str, name2: str) -> Optional[Dict]:
        """
        获取两个标记的相对位置（线程安全）
        
        Args:
            name1 (str): 第一个标记名称
            name2 (str): 第二个标记名称
            
        Returns:
            Dict or None: 相对位置信息，如果任一标记不存在或未跟踪返回None
        """
        with self._lock:
            # 检查标记是否存在
            if name1 not in self.registered_markers:
                print(f"标记 '{name1}' 不存在")
                return None
            
            if name2 not in self.registered_markers:
                print(f"标记 '{name2}' 不存在")
                return None
            
            marker1 = self.registered_markers[name1]
            marker2 = self.registered_markers[name2]
            
            # 检查标记是否被跟踪
            if not marker1.is_tracked:
                print(f"标记 '{name1}' 当前未被跟踪")
                return None
            
            if not marker2.is_tracked:
                print(f"标记 '{name2}' 当前未被跟踪")
                return None
            
            # 计算相对位置
            relative_position = marker2.position - marker1.position
            relative_distance = np.linalg.norm(relative_position)
            
            # 计算相对姿态
            relative_rotation = marker2.rotation - marker1.rotation
            
            # 处理角度环绕
            relative_rotation = np.where(relative_rotation > 180, relative_rotation - 360, relative_rotation)
            relative_rotation = np.where(relative_rotation < -180, relative_rotation + 360, relative_rotation)
            
            # 计算变换矩阵
            T1 = marker1.get_transform_matrix()
            T2 = marker2.get_transform_matrix()
            
            if T1 is not None and T2 is not None:
                # 从marker1到marker2的变换
                T_relative = np.linalg.inv(T1) @ T2
            else:
                T_relative = None
            
            return {
                'marker1': name1,
                'marker2': name2,
                'relative_position': relative_position.tolist(),
                'relative_distance': relative_distance,
                'relative_rotation': relative_rotation.tolist(),
                'transform_matrix': T_relative.tolist() if T_relative is not None else None
            }
    
    def get_all_markers_status(self) -> Dict[str, Dict]:
        """
        获取所有标记的状态（线程安全）
        
        Returns:
            Dict[str, Dict]: 所有标记的状态字典
        """
        with self._lock:
            return {name: marker.get_status() for name, marker in self.registered_markers.items()}
    
    def get_tracking_statistics(self) -> Dict:
        """
        获取跟踪统计信息（线程安全）
        
        Returns:
            Dict: 跟踪统计信息
        """
        with self._lock:
            tracked_count = sum(1 for marker in self.registered_markers.values() if marker.is_tracked)
            total_count = len(self.registered_markers)
            
            return {
                'total_markers': total_count,
                'tracked_markers': tracked_count,
                'lost_markers': total_count - tracked_count,
                'detection_count': self.detection_count,
                'fps': self.fps,
                'is_tracking': self.is_tracking,
                'calibrated': self.calibrated
            }
    
    def save_tracking_data(self, filename: str) -> bool:
        """
        保存跟踪数据
        
        Args:
            filename (str): 保存文件名
            
        Returns:
            bool: 是否成功保存
        """
        try:
            data = {
                'timestamp': time.time(),
                'statistics': self.get_tracking_statistics(),
                'markers': self.get_all_markers_status()
            }
            
            with open(filename, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=4, ensure_ascii=False)
            
            print(f"跟踪数据已保存到: {filename}")
            return True
            
        except Exception as e:
            print(f"保存跟踪数据失败: {e}")
            return False
    
    def is_tracking_active(self) -> bool:
        """
        检查跟踪是否正在运行
        
        Returns:
            bool: 跟踪是否正在运行
        """
        return self.is_tracking and (self.tracking_thread is not None and self.tracking_thread.is_alive())
    
    def wait_for_tracking_start(self, timeout: float = 5.0) -> bool:
        """
        等待跟踪开始
        
        Args:
            timeout (float): 等待超时时间（秒）
            
        Returns:
            bool: 是否成功开始跟踪
        """
        start_time = time.time()
        while time.time() - start_time < timeout:
            if self.is_tracking_active():
                return True
            time.sleep(0.1)
        return False
    
    def get_current_frame(self) -> Optional[np.ndarray]:
        """
        获取当前帧（线程安全）
        
        Returns:
            Optional[np.ndarray]: 当前帧图像，如果没有则返回None
        """
        if self.camera is not None:
            return self.camera.frame
        return None
    
    def __enter__(self):
        """上下文管理器入口"""
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器出口"""
        self.stop_tracking()
    
    def __del__(self):
        """析构函数"""
        if self.is_tracking:
            self.stop_tracking()
